/* Based on the public domain implementation in
 * crypto_hash/keccakc512/simple/ from http://bench.cr.yp.to/supercop.html
 * by Ronny Van Keer
 * and the public domain "TweetFips202" implementation
 * from https://twitter.com/tweetfips202
 * by Gilles Van Assche, Daniel J. Bernstein, and Peter Schwabe */

#include <stdint.h>
#include <assert.h>
#include "fips202.h"

#define NROUNDS 24
#define ROL(a, offset) ((a << offset) ^ (a >> (64-offset)))

/*************************************************
* Name:        load64
*
* Description: Load 8 bytes into uint64_t in little-endian order
*
* Arguments:   - const unsigned char *x: pointer to input byte array
*
* Returns the loaded 64-bit unsigned integer
**************************************************/
static uint64_t load64(const unsigned char *x)
{
  unsigned long long r = 0, i;

  for (i = 0; i < 8; ++i) {
    r |= (unsigned long long)x[i] << 8 * i;
  }
  return r;
}

/*************************************************
* Name:        store64
*
* Description: Store a 64-bit integer to a byte array in little-endian order
*
* Arguments:   - uint8_t *x: pointer to the output byte array
*              - uint64_t u: input 64-bit unsigned integer
**************************************************/
static void store64(uint8_t *x, uint64_t u)
{
  unsigned int i;

  for(i=0; i<8; ++i) {
    x[i] = u;
    u >>= 8;
  }
}

#include <string.h>
#define MIN(a, b) ((a) < (b) ? (a) : (b))


/*************************************************
* Name:        keccak_absorb
*
* Description: Absorb step of Keccak;
*              non-incremental, starts by zeroeing the state.
*
* Arguments:   - uint64_t *s:             pointer to (uninitialized) output Keccak state
*              - unsigned int r:          rate in bytes (e.g., 168 for SHAKE128)
*              - const unsigned char *m:  pointer to input to be absorbed into s
*              - unsigned long long mlen: length of input in bytes
*              - unsigned char p:         domain-separation byte for different Keccak-derived functions
**************************************************/
void keccak_absorb(uint64_t *s,
                          unsigned int r,
                          const unsigned char *m, unsigned long long int mlen,
                          unsigned char p)
{
  unsigned long long i;
  unsigned char t[200];

  // Zero state
  for (i = 0; i < 25; ++i)
    s[i] = 0;

  while (mlen >= r)
  {
    for (i = 0; i < r / 8; ++i)
      s[i] ^= load64(m + 8 * i);

    KeccakF1600_StatePermute(s);
    mlen -= r;
    m += r;
  }

  for (i = 0; i < r; ++i)
    t[i] = 0;
  for (i = 0; i < mlen; ++i)
    t[i] = m[i];
  t[i] = p;
  t[r - 1] |= 128;
  for (i = 0; i < r / 8; ++i)
    s[i] ^= load64(t + 8 * i);
}


/*************************************************
* Name:        keccak_squeezeblocks
*
* Description: Squeeze step of Keccak. Squeezes full blocks of r bytes each.
*              Modifies the state. Can be called multiple times to keep squeezing,
*              i.e., is incremental.
*
* Arguments:   - unsigned char *h:               pointer to output blocks
*              - unsigned long long int nblocks: number of blocks to be squeezed (written to h)
*              - uint64_t *s:                    pointer to in/output Keccak state
*              - unsigned int r:                 rate in bytes (e.g., 168 for SHAKE128)
**************************************************/
void keccak_squeezeblocks(unsigned char *h, unsigned long long int nblocks,
                                 uint64_t *s,
                                 unsigned int r)
{
  unsigned int i;
  while(nblocks > 0)
  {
    KeccakF1600_StatePermute(s);
    for(i=0;i<(r>>3);i++)
    {
      store64(h+8*i, s[i]);
    }
    h += r;
    nblocks--;
  }
}

/********** cSHAKE128 ***********/

// void cshake128_simple_absorb(uint64_t s[25], uint16_t cstm, const unsigned char *in, unsigned long long inlen)
// {
//   unsigned char *sep = (unsigned char*)s;
//   unsigned int i;

//   for (i = 0; i < 25; i++)
//     s[i] = 0;

//   /* Absorb customization (domain-separation) string */
//   sep[0] = 0x01;
//   sep[1] = 0xa8;
//   sep[2] = 0x01;
//   sep[3] = 0x00;
//   sep[4] = 0x01;
//   sep[5] = 16; // fixed bitlen of cstm
//   sep[6] = cstm & 0xff;
//   sep[7] = cstm >> 8;

//   KeccakF1600_StatePermute(s);

//   /* Absorb input */
//   keccak_absorb(s, SHAKE128_RATE, in, inlen, 0x04);
// }


// void cshake128_simple_squeezeblocks(unsigned char *output, unsigned long long nblocks, uint64_t *s)
// {
//   keccak_squeezeblocks(output, nblocks, s, SHAKE128_RATE);
// }


// void cshake128_simple(unsigned char *output, unsigned long long outlen, uint16_t cstm, const unsigned char *in, unsigned long long inlen)
// {
//   uint64_t s[25];
//   unsigned char t[SHAKE128_RATE];
//   unsigned int i;

//   cshake128_simple_absorb(s, cstm, in, inlen);

//   /* Squeeze output */
//   keccak_squeezeblocks(output, outlen/SHAKE128_RATE, s, SHAKE128_RATE);
//   output += (outlen/SHAKE128_RATE)*SHAKE128_RATE;

//   if (outlen%SHAKE128_RATE)
//   {
//     keccak_squeezeblocks(t, 1, s, SHAKE128_RATE);
//     for (i = 0; i < outlen%SHAKE128_RATE; i++)
//       output[i] = t[i];
//   }
// }



// /*************************************************
// * Name:        shake128_absorb
// *
// * Description: Absorb step of the SHAKE128 XOF.
// *              non-incremental, starts by zeroeing the state.
// *
// * Arguments:   - uint64_t *s:                     pointer to (uninitialized) output Keccak state
// *              - const unsigned char *input:      pointer to input to be absorbed into s
// *              - unsigned long long inputByteLen: length of input in bytes
// **************************************************/
// void shake128_absorb(uint64_t *s, const unsigned char *input, unsigned int inputByteLen)
// {
//   keccak_absorb(s, SHAKE128_RATE, input, inputByteLen, 0x1F);
// }

// /*************************************************
// * Name:        shake128_squeezeblocks
// *
// * Description: Squeeze step of SHAKE128 XOF. Squeezes full blocks of SHAKE128_RATE bytes each.
// *              Modifies the state. Can be called multiple times to keep squeezing,
// *              i.e., is incremental.
// *
// * Arguments:   - unsigned char *output:      pointer to output blocks
// *              - unsigned long long nblocks: number of blocks to be squeezed (written to output)
// *              - uint64_t *s:                pointer to in/output Keccak state
// **************************************************/
// void shake128_squeezeblocks(unsigned char *output, unsigned long long nblocks, uint64_t *s)
// {
//   keccak_squeezeblocks(output, nblocks, s, SHAKE128_RATE);
// }

// void shake128(unsigned char *output, unsigned long long outlen, const unsigned char *input,  unsigned long long inlen)
// {
//   uint64_t s[25] = {0};
//   unsigned char t[SHAKE128_RATE];
//   unsigned long long nblocks = outlen/SHAKE128_RATE;
//   size_t i;

//   /* Absorb input */
//   keccak_absorb(s, SHAKE128_RATE, input, inlen, 0x1F);

//   /* Squeeze output */
//   keccak_squeezeblocks(output, nblocks, s, SHAKE128_RATE);

//   output += nblocks*SHAKE128_RATE;
//   outlen -= nblocks*SHAKE128_RATE;

//   if (outlen)
//   {
//     keccak_squeezeblocks(t, 1, s, SHAKE128_RATE);
//     for (i = 0; i < outlen; i++)
//       output[i] = t[i];
//   }
// }


// void shake256_absorb(uint64_t *s, const unsigned char *input, unsigned int inputByteLen)
// {
// 	keccak_absorb(s, SHAKE256_RATE, input, inputByteLen, 0x1F);
// }


// void shake256_squeezeblocks(unsigned char *output, unsigned long long nblocks, uint64_t *s)
// {
// 	keccak_squeezeblocks(output, nblocks, s, SHAKE256_RATE);
// }

// ************************************************
// * Name:        shake256
// *
// * Description: SHAKE256 XOF with non-incremental API
// *
// * Arguments:   - unsigned char *output:      pointer to output
// *              - unsigned long long outlen:  requested output length in bytes
//                - const unsigned char *input: pointer to input
//                - unsigned long long inlen:   length of input in bytes
// *************************************************
// void shake256(unsigned char *output, unsigned long long outlen,
//               const unsigned char *input,  unsigned long long inlen)
// {
//   uint64_t s[25];
//   unsigned char t[SHAKE256_RATE];
//   unsigned long long nblocks = outlen/SHAKE256_RATE;
//   size_t i;

//   /* Absorb input */
//   keccak_absorb(s, SHAKE256_RATE, input, inlen, 0x1F);

//   /* Squeeze output */
//   keccak_squeezeblocks(output, nblocks, s, SHAKE256_RATE);

//   output+=nblocks*SHAKE256_RATE;
//   outlen-=nblocks*SHAKE256_RATE;

//   if(outlen)
//   {
//     keccak_squeezeblocks(t, 1, s, SHAKE256_RATE);
//     for(i=0;i<outlen;i++)
//       output[i] = t[i];
//   }
// }

// /*************************************************
// * Name:        sha3_256
// *
// * Description: SHA3-256 with non-incremental API
// *
// * Arguments:   - unsigned char *output:      pointer to output
// *              - const unsigned char *input: pointer to input
// *              - unsigned long long inlen:   length of input in bytes
// **************************************************/
// void sha3_256(unsigned char *output, const unsigned char *input,  unsigned long long inlen)
// {
//   uint64_t s[25];
//   unsigned char t[SHA3_256_RATE];
//   size_t i;

//   /* Absorb input */
//   keccak_absorb(s, SHA3_256_RATE, input, inlen, 0x06);

//   /* Squeeze output */
//   keccak_squeezeblocks(t, 1, s, SHA3_256_RATE);

//   for(i=0;i<32;i++)
//       output[i] = t[i];
// }

// /*************************************************
// * Name:        sha3_512
// *
// * Description: SHA3-512 with non-incremental API
// *
// * Arguments:   - unsigned char *output:      pointer to output
// *              - const unsigned char *input: pointer to input
// *              - unsigned long long inlen:   length of input in bytes
// **************************************************/
// void sha3_512(unsigned char *output, const unsigned char *input,  unsigned long long inlen)
// {
//   uint64_t s[25];
//   unsigned char t[SHA3_512_RATE];
//   size_t i;

//   /* Absorb input */
//   keccak_absorb(s, SHA3_512_RATE, input, inlen, 0x06);

//   /* Squeeze output */
//   keccak_squeezeblocks(t, 1, s, SHA3_512_RATE);

//   for(i=0;i<64;i++)
//       output[i] = t[i];
// }

// /********** cSHAKE256 ***********/

// void cshake256_simple_absorb(uint64_t s[25], uint16_t cstm, const unsigned char *in, unsigned long long inlen)
// {
//   unsigned char *sep = (unsigned char*)s;
//   unsigned int i;

//   for (i = 0; i < 25; i++)
//     s[i] = 0;

//   /* Absorb customization (domain-separation) string */
//   sep[0] = 0x01;
//   sep[1] = 0x88;
//   sep[2] = 0x01;
//   sep[3] = 0x00;
//   sep[4] = 0x01;
//   sep[5] = 16; // fixed bitlen of cstm
//   sep[6] = cstm & 0xff;
//   sep[7] = cstm >> 8;

//   KeccakF1600_StatePermute(s);

//   /* Absorb input */
//   keccak_absorb(s, SHAKE256_RATE, in, inlen, 0x04);
// }


// void cshake256_simple_squeezeblocks(unsigned char *output, unsigned long long nblocks, uint64_t *s)
// {
//   keccak_squeezeblocks(output, nblocks, s, SHAKE256_RATE);
// }


// void cshake256_simple(unsigned char *output, unsigned long long outlen, uint16_t cstm, const unsigned char *in, unsigned long long inlen)
// {
//   uint64_t s[25];
//   unsigned char t[SHAKE256_RATE];
//   unsigned int i;

//   cshake256_simple_absorb(s, cstm, in, inlen);

//   /* Squeeze output */
//   keccak_squeezeblocks(output, outlen/SHAKE256_RATE, s, SHAKE256_RATE);
//   output += (outlen/SHAKE256_RATE)*SHAKE256_RATE;

//   if(outlen%SHAKE256_RATE)
//   {
//     keccak_squeezeblocks(t, 1, s, SHAKE256_RATE);
//     for (i = 0; i < outlen%SHAKE256_RATE; i++)
//       output[i] = t[i];
//   }
// }
